import socket
import struct
import threading
import time
from typing import Callable, Dict, List, Union

from dataTemp import MesTemp_pb2


class SocketWrapper:
    def __init__(self, host: str, port: int, on_connect_success: Callable[[], None] = None, on_connect_failure: Callable[[Exception], None] = None):
        self.host = host
        self.port = port
        self.on_connect_success = on_connect_success
        self.on_connect_failure = on_connect_failure
        self.conn = None
        self.buf = bytearray(65536)  # 调整缓冲区大小根据需要
        self.subscribers: Dict[int, List[Callable[[MesTemp_pb2.Mes], None]]] = {}
        self._running = False
        self._thread = None

        self._connect()

    def _connect(self) -> None:
        try:
            self.conn = socket.create_connection((self.host, self.port))
            self._running = True
            self._thread = threading.Thread(target=self._receive_loop, daemon=True)
            self._thread.start()
            if self.on_connect_success:
                self.on_connect_success()
        except Exception as e:
            if self.on_connect_failure:
                self.on_connect_failure(e)

    def send_Mes(self, data, mes_type: int) -> None:
        # 序列化 Mes
        data = data.SerializeToString()
        # 创建 Mes 对象并设置字段
        mes = MesTemp_pb2.Mes()
        mes.Type = mes_type
        mes.Data = data
        # 序列化 Mes 对象
        serialized_mes = mes.SerializeToString()
        # 发送数据
        self._send_with_length(serialized_mes)

    def close(self) -> None:
        self._running = False
        if self._thread:
            self._thread.join()
        if self.conn:
            self.conn.close()

    def _send_with_length(self, data: bytes) -> None:
        # 前四个字节为数据长度
        length = len(data)
        header = struct.pack('!I', length)
        self.conn.sendall(header + data)

    def _receive_with_length(self) -> bytes:
        # 读取数据长度
        header = self.conn.recv(4)
        length = struct.unpack('!I', header)[0]
        # 读取实际数据
        data = self.conn.recv(length)
        while len(data) < length:
            data += self.conn.recv(length - len(data))
        return data

    @staticmethod
    def getMsg(obj, tp: int, callback: Callable[[MesTemp_pb2.Mes], None]) -> None:
        # 订阅消息类型，并设置回调
        def Mes_handler(Mes: MesTemp_pb2.Mes) -> None:
            obj.ParseFromString(Mes.Data)
            callback(obj)
        # 实例化 SocketWrapper 对象时添加回调
        # 在使用时需要手动调用 subscribe 方法
        pass

    def _receive_loop(self) -> None:
        while self._running:
            try:
                serialized_mes = self._receive_with_length()
                mes = MesTemp_pb2.Mes()
                mes.ParseFromString(serialized_mes)
                self._notify_subscribers(mes)
            except Exception as e:
                print(f"Error in receive loop: {e}")

    def subscribe(self, mes_type: int, callback: Callable[[MesTemp_pb2.Mes], None]) -> None:
        if mes_type not in self.subscribers:
            self.subscribers[mes_type] = []
        self.subscribers[mes_type].append(callback)

    def unsubscribe(self, mes_type: int, callback: Callable[[MesTemp_pb2.Mes], None]) -> None:
        if mes_type in self.subscribers:
            if callback in self.subscribers[mes_type]:
                self.subscribers[mes_type].remove(callback)
                if not self.subscribers[mes_type]:
                    del self.subscribers[mes_type]

    def _notify_subscribers(self, mes: MesTemp_pb2.Mes) -> None:
        if mes.Type in self.subscribers:
            for callback in self.subscribers[mes.Type]:
                try:
                    callback(mes)
                except Exception as e:
                    print(f"Error in callback: {e}")


# 使用示例

def on_connect_success():
    print("Successfully connected to the server.")

def on_connect_failure(e: Exception):
    print(f"Failed to connect to the server: {e}")

def Mes_handler(Mes: MesTemp_pb2.Mes) -> None:
    # 创建一个 User 实例
    user = MesTemp_pb2.User()
    user.ParseFromString(Mes.Data)
    print("Received Mes:", user)

if __name__ == "__main__":
    # 创建连接
    socket_wrapper = SocketWrapper('192.168.1.7', 8000, on_connect_success, on_connect_failure)

    # 创建一个 User 实例
    user = MesTemp_pb2.User()
    user.pwd = "123"

    # 订阅消息类型
    socket_wrapper.subscribe(MesTemp_pb2.LoginMT, Mes_handler)

    # 发送消息
    socket_wrapper.send_Mes(user, MesTemp_pb2.LoginMT)
    time.sleep(10)
    # 关闭连接
    socket_wrapper.close()
